# Copyright 2020-2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Res2net backbone."""

import math
import numpy as np
import mindspore.nn as nn
from mindspore.ops import operations as P
from mindspore.common.tensor import Tensor
from mindspore.ops import functional as F


def weight_init_ones(shape):
    """Weight init."""
    return Tensor(np.full(shape, 0.01).astype(np.float32))


def _conv(
        in_channels, out_channels, kernel_size=3, stride=1, padding=0, pad_mode="pad"
):
    """Conv2D wrapper."""
    shape = (out_channels, in_channels, kernel_size, kernel_size)
    weights = weight_init_ones(shape)
    return nn.Conv2d(
        in_channels,
        out_channels,
        kernel_size=kernel_size,
        stride=stride,
        padding=padding,
        pad_mode=pad_mode,
        weight_init=weights,
        has_bias=False,
    )


def _BatchNorm2dInit(out_chls, momentum=0.1, affine=True, use_batch_statistics=True):
    """Batchnorm2D wrapper."""
    dtype = np.float32
    gamma_init = Tensor(np.array(np.ones(out_chls)).astype(dtype))
    beta_init = Tensor(np.array(np.ones(out_chls) * 0).astype(dtype))
    moving_mean_init = Tensor(np.array(np.ones(out_chls) * 0).astype(dtype))
    moving_var_init = Tensor(np.array(np.ones(out_chls)).astype(dtype))
    return nn.BatchNorm2d(
        out_chls,
        momentum=momentum,
        affine=affine,
        gamma_init=gamma_init,
        beta_init=beta_init,
        moving_mean_init=moving_mean_init,
        moving_var_init=moving_var_init,
        use_batch_statistics=use_batch_statistics,
    )


class Res2NetFea(nn.Cell):
    """
    Res2Net architecture.

    Args:
        block (Cell): Block for network.
        layer_nums (list): Numbers of block in different layers.
        in_channels (list): Input channel in each layer.
        out_channels (list): Output channel in each layer.
        weights_update (bool): Weight update flag.
    Returns:
        Tensor, output tensor.

    Examples:
        >>> Res2Net(Residual2Block,
        >>>        [3, 4, 6, 3],
        >>>        [64, 256, 512, 1024],
        >>>        [256, 512, 1024, 2048],
        >>>        False)
    """

    def __init__(
            self, block, layer_nums, in_channels, out_channels, weights_update=False, detect=True
    ):
        super(Res2NetFea, self).__init__()

        if not len(layer_nums) == len(in_channels) == len(out_channels) == 4:
            raise ValueError(
                "the length of " "layer_num, inchannel, outchannel list must be 4!"
            )
        self.detect = detect
        bn_training = False
        self.conv1_0 = _conv(3, 32, kernel_size=3, stride=2, padding=0, pad_mode="same")
        self.bn1_0 = _BatchNorm2dInit(32)
        self.conv1_1 = _conv(
            32, 32, kernel_size=3, stride=1, padding=0, pad_mode="same"
        )
        self.bn1_1 = _BatchNorm2dInit(32)
        self.conv1_2 = _conv(
            32,
            64,
            kernel_size=3,
            stride=1,
            padding=0,
            pad_mode="same",
        )
        self.bn1 = _BatchNorm2dInit(
            64, affine=bn_training, use_batch_statistics=bn_training
        )
        self.bn1 = _BatchNorm2dInit(
            64, affine=bn_training, use_batch_statistics=bn_training
        )
        self.relu = P.ReLU()
        self.maxpool = P.MaxPool(kernel_size=3, strides=2, pad_mode="SAME")
        self.weights_update = weights_update

        if not self.weights_update:
            self.conv1_0.weight.requires_grad = False
            self.conv1_1.weight.requires_grad = False
            self.conv1_2.weight.requires_grad = False

        self.layer1 = self._make_layer(
            block,
            layer_nums[0],
            in_channel=in_channels[0],
            out_channel=out_channels[0],
            stride=1,
            training=bn_training,
            weights_update=self.weights_update,
        )
        self.layer2 = self._make_layer(
            block,
            layer_nums[1],
            in_channel=in_channels[1],
            out_channel=out_channels[1],
            stride=2,
            training=bn_training,
            weights_update=True,
        )
        self.layer3 = self._make_layer(
            block,
            layer_nums[2],
            in_channel=in_channels[2],
            out_channel=out_channels[2],
            stride=2,
            training=bn_training,
            weights_update=True,
        )
        self.layer4 = self._make_layer(
            block,
            layer_nums[3],
            in_channel=in_channels[3],
            out_channel=out_channels[3],
            stride=2,
            training=bn_training,
            weights_update=True,
        )

    def _make_layer(
            self,
            block,
            layer_num,
            in_channel,
            out_channel,
            stride,
            training=False,
            weights_update=False,
    ):
        """Make block layer."""
        layers = []
        down_sample = False
        if stride != 1 or in_channel != out_channel:
            down_sample = True
        resblk = block(
            in_channel,
            out_channel,
            stride=stride,
            down_sample=down_sample,
            training=training,
            weights_update=weights_update,
            stype="stage",
        )
        layers.append(resblk)

        for _ in range(1, layer_num):
            resblk = block(
                out_channel,
                out_channel,
                stride=1,
                training=training,
                weights_update=weights_update,
            )
            layers.append(resblk)

        return nn.SequentialCell(layers)

    def construct(self, x):
        """
        construct the Res2Net Network

        Args:
            x: input feature data.

        Returns:
        Tensor, output tensor.
        """
        x = self.conv1_0(x)
        x = self.bn1_0(x)
        x = self.relu(x)
        x = self.conv1_1(x)
        x = self.bn1_1(x)
        x = self.relu(x)
        x = self.conv1_2(x)
        x = self.bn1(x)
        x = self.relu(x)
        c1 = self.maxpool(x)

        c2 = self.layer1(c1)
        identity = c2
        if not self.weights_update:
            identity = F.stop_gradient(c2)
        c3 = self.layer2(identity)
        c4 = self.layer3(c3)
        c5 = self.layer4(c4)
        if self.detect:
            return c3, c4, c5
        return c5


class Residual2BlockUsing(nn.Cell):
    """
    Res2Net residual block definition.

    Args:
        in_channels (int) - Input channel.
        out_channels (int) - Output channel.
        stride (int) - Stride size for the initial convolutional layer. Default: 1.
        down_sample (bool) - If to do the downsample in block. Default: False.
        momentum (float) - Momentum for batchnorm layer. Default: 0.1.
        training (bool) - Training flag. Default: False.
        weights_updata (bool) - Weights update flag. Default: False.

    Returns:
        Tensor, output tensor.

    Examples:
        ResidualBlock(3,256,stride=2,down_sample=True)
    """

    expansion = 4

    def __init__(
            self,
            in_channels,
            out_channels,
            stride=1,
            down_sample=False,
            momentum=0.1,
            training=False,
            weights_update=False,
            baseWidth=26,
            scale=4,
            stype="normal",
    ):
        super(Residual2BlockUsing, self).__init__()

        self.affine = weights_update

        # out_chls = out_channels // self.expansion
        assert scale > 1, "Res2Net is ResNet when scale = 1"
        width = int(math.floor(out_channels // self.expansion * (baseWidth / 64.0)))
        channel = width * scale
        self.conv1 = _conv(in_channels, channel, kernel_size=1, stride=1, padding=0)
        self.bn1 = _BatchNorm2dInit(
            channel,
            momentum=momentum,
            affine=self.affine,
            use_batch_statistics=training,
        )
        if stype == "stage":
            self.pool = nn.AvgPool2d(kernel_size=3, stride=stride, pad_mode="same")

        self.convs = nn.CellList()
        self.bns = nn.CellList()
        for _ in range(scale - 1):
            self.convs.append(
                _conv(width, width, kernel_size=3, stride=stride, padding=1)
            )
            self.bns.append(
                _BatchNorm2dInit(
                    width,
                    momentum=momentum,
                    affine=self.affine,
                    use_batch_statistics=training,
                )
            )

        self.conv3 = _conv(channel, out_channels, kernel_size=1, stride=1, padding=0)
        self.bn3 = _BatchNorm2dInit(
            out_channels,
            momentum=momentum,
            affine=self.affine,
            use_batch_statistics=training,
        )

        if training:
            self.bn1 = self.bn1.set_train()
            for each_bn in self.bns:
                each_bn.set_train()
            self.bn3 = self.bn3.set_train()

        if not weights_update:
            self.conv1.weight.requires_grad = False
            for each_conv in self.convs:
                each_conv.weight.requires_grad = False
            self.conv3.weight.requires_grad = False

        self.relu = P.ReLU()
        self.downsample = down_sample
        if self.downsample:
            if stride == 1:
                self.conv_down_sample = _conv(
                    in_channels, out_channels, kernel_size=1, stride=1, padding=0
                )
                self.bn_down_sample = _BatchNorm2dInit(
                    out_channels,
                    momentum=momentum,
                    affine=self.affine,
                    use_batch_statistics=training,
                )
            else:
                self.pool_down_sample = nn.MaxPool2d(
                    kernel_size=2, stride=2, pad_mode="same"
                )
                self.conv_down_sample = _conv(
                    in_channels, out_channels, kernel_size=1, stride=1, padding=0
                )
                self.bn_down_sample = _BatchNorm2dInit(
                    out_channels,
                    momentum=momentum,
                    affine=self.affine,
                    use_batch_statistics=training,
                )
            if training:
                self.bn_down_sample = self.bn_down_sample.set_train()
            if not weights_update:
                self.conv_down_sample.weight.requires_grad = False
        self.add = P.Add()
        self.scale = scale
        self.width = width
        self.stride = stride
        self.stype = stype
        self.split = P.Split(axis=1, output_num=scale)
        self.cat = P.Concat(axis=1)

    def construct(self, x):
        """
        construct the Res2Net residual block

        Args:
            x: input feature data.

        Returns:
        Tensor, output tensor.
        """
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        spx = self.split(out)

        sp = self.convs[0](spx[0])
        sp = self.relu(self.bns[0](sp))
        out = sp
        for i in range(1, self.scale - 1):
            if self.stype == "stage":
                sp = spx[i]
            else:
                # sp = sp[:, :, :, :]  # to avoid bug in mindspore
                sp = sp + spx[i]
            sp = self.convs[i](sp)
            sp = self.relu(self.bns[i](sp))
            out = self.cat((out, sp))

        if self.stype == "normal":
            out = self.cat((out, spx[self.scale - 1]))
        elif self.stype == "stage":
            out = self.cat((out, self.pool(spx[self.scale - 1])))

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample:
            if self.stride != 1:
                identity = self.pool_down_sample(identity)
            identity = self.conv_down_sample(identity)
            identity = self.bn_down_sample(identity)

        out = self.add(out, identity)
        out = self.relu(out)

        return out
